// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/testing/expression/spark_special_form_signature_generator.h>

namespace kumo::pollux::fuzzer {

std::vector<exec::FunctionSignaturePtr>
SparkSpecialFormSignatureGenerator::getSignaturesForCast() const {
  std::vector<exec::FunctionSignaturePtr> signatures =
      SpecialFormSignatureGenerator::getSignaturesForCast();

  // Cast tinyint/smallint/integer/bigint as varbinary is supported in Spark.
  for (auto fromType : {"tinyint", "smallint", "integer", "bigint"}) {
    signatures.push_back(makeCastSignature(fromType, "varbinary"));
  }

  // Cast tinyint/smallint/integer/bigint as timestamp is supported in Spark.
  for (auto fromType : {"tinyint", "smallint", "integer", "bigint"}) {
    signatures.push_back(makeCastSignature(fromType, "timestamp"));
  }
  return signatures;
}

const std::unordered_map<std::string, std::vector<exec::FunctionSignaturePtr>>&
SparkSpecialFormSignatureGenerator::getSignatures() const {
  const static std::
      unordered_map<std::string, std::vector<exec::FunctionSignaturePtr>>
          kSpecialForms{
              {"and", getSignaturesForAnd()},
              {"or", getSignaturesForOr()},
              {"coalesce", getSignaturesForCoalesce()},
              {"if", getSignaturesForIf()},
              {"switch", getSignaturesForSwitch()},
              {"cast", getSignaturesForCast()},
              {"concat_ws", getSignaturesForConcatWs()}};
  return kSpecialForms;
}

std::vector<exec::FunctionSignaturePtr>
SparkSpecialFormSignatureGenerator::getSignaturesForConcatWs() const {
  // Signature: concat_ws (separator, input, ...) -> output:
  // varchar, varchar, varchar, ... -> varchar
  return {kumo::pollux::exec::FunctionSignatureBuilder()
              .argumentType("varchar")
              .variableArity("varchar")
              .returnType("varchar")
              .build()};
}

} // namespace kumo::pollux::fuzzer
